import fitz  # PyMuPDF
from utils.TextChunker import TextChunker

class PDFTextExtractor:
    def __init__(self, pdf_path):
        """
        初始化 PDF 文本提取器
        :param pdf_path: PDF 文件路径
        """
        self.pdf_path = pdf_path
        self.doc = fitz.open(pdf_path)
        self.page_count = self.doc.page_count

    def get_all_pages_text(self):
        """
        获取所有页的文本列表
        :return: dict[str,str], 每页的文本内容
        """
        text_list = []
        for page_num in range(self.page_count):
            page = self.doc.load_page(page_num)
            text = page.get_text()
            text_list.append(text) 
        return text_list

    def get_merged_pages_text(self):
        """
        获取合并关联后的页文本字典
        :return: dict, 格式如 {"1,2,3":"文本1+文本2+文本3"}
        """
        merged_dict = {}
        all_texts = self.get_all_pages_text()
        chunks = TextChunker().chunk_text(''.join(all_texts))
        for index , chunk in enumerate(chunks):
            merged_dict[str(index+1)] = chunk
        return merged_dict

    def close(self):
        """关闭 PDF 文档"""
        self.doc.close()

    def __enter__(self):
        """支持上下文管理"""
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """支持上下文管理"""
        self.close()


# 使用示例
if __name__ == "__main__":
    pdf_path = "美国儿科学会育儿百科(第7版).pdf"  # 替换为你的PDF文件路径

    with PDFTextExtractor(pdf_path) as extractor:
        # 获取所有页文本
        all_pages = extractor.get_all_pages_text()
        print(f"总页数: {len(all_pages)}")
        print("第一页内容:", all_pages['1'][:100] + "...")  # 打印第一页前100个字符

        # 获取合并关联页文本
        merged_pages = extractor.get_merged_pages_text()
        print("\n合并关联页示例:")
        for key, value in list(merged_pages.items())[100:130]:  # 打印前3个合并项
            print(f"页 {key}: {value[:100]}...")